package com.tianji.aigc.memory;


import cn.hutool.core.bean.BeanUtil;
import cn.hutool.json.JSONUtil;
import com.tianji.aigc.entity.RedisMessage;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.util.List;

/**
 * 基于redis的会话存储类
 */
@Component
public class RedisChatMemory implements ChatMemory {


    @Autowired
    private StringRedisTemplate redisTemplate;

    /**
     * 对话完成添加消息
     *
     * @param conversationId
     * @param messages
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        //1. 拼接RedisKey
        String redisKey = getRedisKey(conversationId);

        //2. 将消息转成字符换类型
        List<String> messageStrList = messages.stream().map(x -> {
            RedisMessage redisMessage = BeanUtil.toBean(x, RedisMessage.class);
            redisMessage.setTextContent(x.getText());
            if (x instanceof AssistantMessage assistantMessage) {
                redisMessage.setToolCalls(assistantMessage.getToolCalls());
            } else if (x instanceof ToolResponseMessage toolResponseMessage) {
                redisMessage.setToolResponses(toolResponseMessage.getResponses());
            }
            return JSONUtil.toJsonStr(redisMessage);
        }).toList();

        //3. 将消息存入缓存中
        redisTemplate.opsForList().rightPushAll(redisKey, messageStrList);
    }

    /**
     * 获取对话ID相关的消息列表
     *
     * @param conversationId
     * @param lastN
     * @return
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        //1. 拼接RedisKey
        String redisKey = getRedisKey(conversationId);

        //2. 获取缓存中的数据
        List<String> messageStrList = redisTemplate.opsForList().range(redisKey, 0, lastN);

        //3. 转成消息对象
        List<Message> messageList = messageStrList.stream().map(x -> {
            RedisMessage redisMessage = JSONUtil.toBean(x, RedisMessage.class);

            Message message = null;
            if (redisMessage.getMessageType().equals(MessageType.ASSISTANT)) {
                message = new AssistantMessage(redisMessage.getTextContent(), redisMessage.getMetadata(), redisMessage.getToolCalls(), redisMessage.getMedia());
            } else if (redisMessage.getMessageType().equals(MessageType.TOOL)) {
                message = new ToolResponseMessage(redisMessage.getToolResponses(), redisMessage.getMetadata());
            } else if (redisMessage.getMessageType().equals(MessageType.USER)) {
                message = new UserMessage(redisMessage.getMessageType(), redisMessage.getTextContent(), redisMessage.getMedia(), redisMessage.getMetadata());
            } else {
                message = new AssistantMessage(redisMessage.getTextContent());
            }
            return message;
        }).toList();

        return messageList;
    }

    /**
     * 清理对话消息
     *
     * @param conversationId
     */
    @Override
    public void clear(String conversationId) {
        //1. 拼接RedisKey
        String redisKey = getRedisKey(conversationId);

        //2. 删除缓存中的数据
        redisTemplate.delete(redisKey);
    }

    /**
     * 拼接RedisKey
     *
     * @param conversationId
     * @return
     */
    private String getRedisKey(String conversationId) {
        return "CHAT:" + conversationId;
    }

}
